import torch
import torch.nn as nn


# use LSTM
class DiscriminatorCandidate2(nn.Module):
    def __init__(self, input_dim, hidden_dim, out_dim, dropout=0.25):
        super(DiscriminatorCandidate2, self).__init__()
        # self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers=2)
        self.rnn = nn.RNN(input_dim, hidden_dim, num_layers=2)
        self.drop_out = nn.Dropout(dropout)
        self.out = nn.Linear(hidden_dim, out_dim)

    def forward(self, input):
        concat_input = input
        _, lstm_out = self.rnn(concat_input.view(concat_input.shape[1], concat_input.shape[0], -1))
        ct = lstm_out[1]
        last_layer_out = ct[-1]
        last_layer_out = last_layer_out.view(last_layer_out.shape[0], 1, -1)

        out = torch.sigmoid(self.out(self.drop_out(last_layer_out)))
        out = out.view(out.shape[0], -1)
        return out
